from toolz.curried import get

import paddle
from paddle.fluid import layers as L

from helm.dynamic.engine.metrics import Average


class TopKAccuracy(Average):

    def __init__(self, name=None, k=5):
        self.k = k
        name = self._init_name(name)
        super().__init__(name, self.output_transform)

    def output_transform(self, output):
        y_pred, y_true = get(["y_pred", "y_true"], output)
        with paddle.fluid.dygraph.no_grad():
            # y_pred = F.softmax(y_pred)
            acc = L.accuracy(y_pred, y_true, k=self.k)
        acc = float(acc.numpy())
        return acc, y_true.shape[0]

    def _init_name(self, name):
        if name is None:
            name = "acc" + str(self.k)
        return name


class Accuracy(Average):

    def __init__(self, name="acc"):
        super().__init__(name, self.output_transform)

    @staticmethod
    def output_transform(output):
        y_pred, y_true = get(["y_pred", "y_true"], output)
        with paddle.fluid.dygraph.no_grad():
            y_pred = L.softmax(y_pred)
            acc = L.accuracy(y_pred, y_true)
        acc = float(acc.numpy())
        return acc, y_true.shape[0]